import pandas as pd
import matplotlib.pyplot as plt
import os


def plot_navs(frequency, begin_date, M=None, N=None):
    """绘制净值曲线"""
    if M is None and N is None:
        files = os.listdir(f'{frequency}/{begin_date}/nav')
    else:
        files = [f'nav_{int(M)}_{float(N)}.csv']
    for file in files:
        print(file)
        nav_df = pd.read_csv(f'{frequency}/{begin_date}/nav/{file}')
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

        # 净值曲线
        ax1.plot(nav_df['date'], nav_df['nav0'], label=f'nav', color='purple')
        ax1.plot(nav_df['date'], nav_df['nav2'], label=f'nav_after_fee', color='blue')
        ax1.plot(nav_df['date'], nav_df['nav(hold_btc)'], label='Hold BTC', color='red')
        ax1.set_title('Strategy Performance')
        ax1.set_ylabel('NAV')
        ax1.set_xlabel('Date')
        ax1.legend()

        ax2.plot(nav_df['date'], nav_df['nav0'], label=f'nav', color='purple')
        ax2.plot(nav_df['date'], nav_df['nav2'], label=f'nav_after_fee', color='blue')
        ax2.plot(nav_df['date'], nav_df['nav(hold_btc)'], label='Hold BTC', color='red')
        ax2.set_yscale('log')  # 设置 y 轴为对数尺度
        ax2.set_title('Strategy Performance (Logarithmic Scale)')
        ax2.set_ylabel('NAV(log)')
        ax2.set_xlabel('Date')
        ax2.legend()

        plt.tight_layout()
        plt.savefig(f'{frequency}/{begin_date}/png/{file[:-4]}.png')
        plt.close(fig)  # 每次循环后关闭


if __name__ == '__main__':
    # 指定参数画净值图
    # plot_navs(frequency='1d', begin_date='20170101', M=7, N=2.5)
    # 画全部净值图
    for frequency in ["1d", "3h", "1h"]:
        for begin_date in ["20170101", "20200101", "20230101"]:
            print(frequency, begin_date)
            plot_navs(frequency, begin_date)